import matplotlib.pyplot as plt
from PIL import Image
import numpy as np
import torch.optim as optim
import os
import slangpy as spy
import pathlib
import random
import math


MAX_NUM_INSTANCES = 8 * 1024 * 1024 # a maximum of 8M stroke instances
NUM_STROKES = 64 * 32

device = spy.create_device(
    type=spy.DeviceType.vulkan, # dx12 does not have support for fp32 atomic accumulation
    include_paths=[
        pathlib.Path(__file__).parent.absolute(),
    ],
    enable_debug_layers=True,
    enable_cuda_interop=True,
    enable_print=True)

class Strokes:
    def __init__(self, num_strokes=0, max_num_instances=MAX_NUM_INSTANCES):
        self.num_strokes = num_strokes
        self.max_num_instances = max_num_instances
        self.stroke_blob = device.create_buffer(size=self.num_strokes * 4 * 6)
        self.d_stroke_blob = device.create_buffer(size=self.num_strokes * 4 * 6)
    def random_initialization(self, width, height):
        np_stroke_positions = np.random.random(size=(self.num_strokes, 2)).astype(np.float32)
        np_stroke_positions[:, 0] *= width
        np_stroke_positions[:, 1] *= height
        num_pixels = width * height
        mean_num_pixels_per_stroke = num_pixels / self.num_strokes
        radius = math.sqrt(mean_num_pixels_per_stroke / (2 * math.pi))
        np_stroke_sizes = np.ones(shape=(self.num_strokes, 1), dtype=np.float32) * radius * 2
        np_stroke_colors = np.random.random(size=(self.num_strokes, 3)).astype(np.float32)
        np_all = np.concatenate([np_stroke_positions, np_stroke_sizes, np_stroke_colors], axis=1)
        self.stroke_blob.copy_from_numpy(np_all)
        # Reset gradient buffer after initialization
        zeros = np.zeros(self.num_strokes * 6, dtype=np.float32)
        self.d_stroke_blob.copy_from_numpy(zeros)

def load_target_image(image_path):
    '''
    Load the target image and convert it to a torch.Tensor
    @param image_path: Path to the target image
    @param image_size: Size to which the image will be resized (default is 256)
    @return: A numpy array of the image with shape (image_size, image_size, 4)
    '''
    width, height = 0, 0
    if os.path.exists(image_path):
        image = Image.open(image_path)
        image_np = np.array(image) / 255.0
        width, height = image.size
        if image_np.shape[2] == 3:
            alpha = np.ones((height, width, 1))
            image_np = np.concatenate([image_np, alpha], axis=2)
            
        return width, height, image_np
    else:
        print(f"Target image {image_path} does not exist.")
        return width, height, np.zeros((height, width, 4), dtype=np.float32)



class Painter:

    def reinit(self):
        self.strokes.random_initialization(self.image_width, self.image_height)

    def setup_uniforms(self):
        ub = {
            'film_dimensions': (self.image_width, self.image_height),
            'tile_dimensions': (self.image_width // 8, self.image_height // 8),
            'standard_stroke_size': self.standard_stroke_size.value,
            'stroke_size_var': self.stroke_size_var.value,
            'num_strokes': self.strokes.num_strokes,
            'max_num_instances': self.strokes.max_num_instances,
            'adam_eta': self.adam_eta.value,
            'lambda_overlap': self.lambda_overlap.value
        }
        return ub


    def __init__(self):

        width, height, np_target_texture = load_target_image(
            image_path="fluffy.png",
        )

        self.image_width = width
        self.image_height = height
        self.strokes = Strokes(num_strokes=NUM_STROKES, max_num_instances=MAX_NUM_INSTANCES)
        self.reinit()

        tile_dimensions = (self.image_width + 7) // 8, (self.image_height + 7) // 8
        num_tiles = tile_dimensions[0] * tile_dimensions[1]
        
        self.instance_stroke_index = device.create_buffer(size=MAX_NUM_INSTANCES * 4, usage=spy.BufferUsage.unordered_access)
        self.instance_count = device.create_buffer(size=4, usage=spy.BufferUsage.unordered_access)
        self.tile_instance_count = device.create_buffer(size=num_tiles * 4, usage=spy.BufferUsage.unordered_access)
        self.tile_instance_offset = device.create_buffer(size=num_tiles * 4, usage=spy.BufferUsage.unordered_access)
        self.tile_instance_count_allocator = device.create_buffer(size=num_tiles * 4, usage=spy.BufferUsage.unordered_access)
        self.adam_first_moment = device.create_buffer(size=self.strokes.num_strokes * 6 * 4, usage=spy.BufferUsage.unordered_access)
        self.adam_second_moment = device.create_buffer(size=self.strokes.num_strokes * 6 * 4, usage=spy.BufferUsage.unordered_access)
        self.output_texture = device.create_texture(format=spy.Format.rgba32_float,
                                                    width=self.image_width,
                                                    height=self.image_height,
                                                    usage=spy.TextureUsage.unordered_access | spy.TextureUsage.copy_source | spy.TextureUsage.shader_resource)
        self.target_texture = device.create_texture(format=spy.Format.rgba32_float,
                                                    width=self.image_width,
                                                    height=self.image_height,
                                                    usage=spy.TextureUsage.copy_destination | spy.TextureUsage.shader_resource)
        
        self.target_texture.copy_from_numpy(np_target_texture.astype(np.float32))

        
        self.sampler = device.create_sampler()

        
        '''
            "ClearCounters",
            "CountInstances",
            "FillInstanceOffsets",
            "SpawnInstances",
            "RenderAndBackward",
            "RenderPerStrokeLossAndBackWard",
            "UpdateParameters",
            "RegularizeParameters"
        '''
        self.ClearCounters_program = device.load_program(
            module_name="slang/stippling.slang",
            entry_point_names=[
            "ClearCounters"
            ]
        )
        self.ClearCounters = device.create_compute_kernel(self.ClearCounters_program)

        self.CountInstances_program = device.load_program(
            module_name="slang/stippling.slang",
            entry_point_names=[
            "CountInstances"
            ]
        )
        self.CountInstances = device.create_compute_kernel(self.CountInstances_program)

        self.FillInstanceOffsets_program = device.load_program(
            module_name="slang/stippling.slang",
            entry_point_names=[
            "FillInstanceOffsets"
            ]
        )
        self.FillInstanceOffsets = device.create_compute_kernel(self.FillInstanceOffsets_program)

        self.SpawnInstances_program = device.load_program(
            module_name="slang/stippling.slang",
            entry_point_names=[
                    "SpawnInstances"
            ]
        )
        self.SpawnInstances = device.create_compute_kernel(self.SpawnInstances_program)

        self.RenderAndBackward_program = device.load_program(
            module_name="slang/stippling.slang",
            entry_point_names=[
            "RenderAndBackward"
            ]
        )
        self.RenderAndBackward = device.create_compute_kernel(self.RenderAndBackward_program)

        self.RenderPerStrokeLossAndBackWard_program = device.load_program(
            module_name="slang/stippling.slang",
            entry_point_names=[
            "RenderPerStrokeLossAndBackWard"
            ]
        )
        self.RenderPerStrokeLossAndBackWard = device.create_compute_kernel(self.RenderPerStrokeLossAndBackWard_program)

        self.UpdateParameters_program = device.load_program(
            module_name="slang/stippling.slang",
            entry_point_names=[
            "UpdateParameters"
            ]
        )
        self.UpdateParameters = device.create_compute_kernel(self.UpdateParameters_program)

        self.RegularizeParameters_program = device.load_program(
            module_name="slang/stippling.slang",
            entry_point_names=["RegularizeParameters"]
        )
        self.RegularizeParameters = device.create_compute_kernel(self.RegularizeParameters_program)


        self.window = spy.Window(width=self.image_width, height=self.image_height,
                             title="Optimizer", resizable=False)
        self.surface = device.create_surface(self.window)
        self.surface.configure(
            format=spy.Format.rgba8_unorm_srgb,
            width=self.window.width,
            height=self.window.height,
            vsync=False
        )

        self.window.on_keyboard_event = self.on_keyboard_event
        self.window.on_mouse_event = self.on_mouse_event
        
        self.ui = spy.ui.Context(device)
        self.setup_ui()
        
        self.uniforms = self.setup_uniforms()
    
    def main(self):
        frame = 0
        timer = spy.Timer()

        random_index = 0
        
        while not self.window.should_close():
            dt = timer.elapsed_s()

            self.fps_text.text = f"FPS: {1.0 / dt:.2f} | Frame: {frame}"

            timer.reset()

            self.window.process_events()
            self.ui.process_events()

            surface_texture = self.surface.acquire_next_image()
            if not surface_texture:
                continue

            command_encoder = device.create_command_encoder()

            self.render(command_encoder)
            
            command_encoder.blit(surface_texture, self.output_texture)

            self.ui.new_frame(surface_texture.width, surface_texture.height)
            self.ui.render(surface_texture, command_encoder)

            device.submit_command_buffer(command_encoder.finish())

            # loss_value = 0
            # self.loss_text.text = f"Loss: {loss_value:.6f} | Random Index: {random_index}"

            del surface_texture

            self.surface.present()

            frame += 1
        
        device.wait()

    def bind(self, pass_encoder, kernel, uniforms):
        """
        Bind the compute kernel and uniforms to the command encoder.
        @param pass_encoder: The command encoder for issuing commands
        @param kernel: The compute kernel to bind
        @param uniforms: The uniform data to bind
        @return: A cursor for binding resources
        """
        shader_obj = pass_encoder.bind_pipeline(kernel)
        cursor = spy.ShaderCursor(shader_obj)
        cursor.ub = uniforms
        cursor.rw_strokes = self.strokes.stroke_blob
        cursor.rw_d_strokes = self.strokes.d_stroke_blob
        cursor.rw_instance_stroke_index = self.instance_stroke_index
        cursor.rw_instance_count = self.instance_count
        cursor.rw_tile_instance_count = self.tile_instance_count
        cursor.rw_tile_instance_offset = self.tile_instance_offset
        cursor.rw_tile_instance_count_allocator = self.tile_instance_count_allocator
        cursor.adam_first_moment = self.adam_first_moment
        cursor.adam_second_moment = self.adam_second_moment
        cursor.rw_output_texture = self.output_texture
        cursor.target_texture = self.target_texture
        cursor.sampler = self.sampler
        return cursor

    def render(self, command_encoder: spy.CommandEncoder):
        """
        Render the scene using the volume primitives and camera.
        @param command_encoder: The command encoder for issuing rendering commands
        """
        self.uniforms = self.setup_uniforms()
        tile_dimensions = (self.image_width + 7) // 8, (self.image_height + 7) // 8
        num_tiles = tile_dimensions[0] * tile_dimensions[1]
        with command_encoder.begin_compute_pass() as pass_encoder:
            pass_encoder.push_debug_group("ClearCounters", (0, 0, 0))
            cursor = self.bind(pass_encoder, self.ClearCounters.pipeline, self.uniforms)
            pass_encoder.dispatch((num_tiles, 1, 1))
            pass_encoder.pop_debug_group()
        with command_encoder.begin_compute_pass() as pass_encoder:
            pass_encoder.push_debug_group("CountInstances", (1, 0, 0))
            cursor = self.bind(pass_encoder, self.CountInstances.pipeline, self.uniforms)
            pass_encoder.dispatch((self.strokes.num_strokes, 1, 1))
            pass_encoder.pop_debug_group()
        with command_encoder.begin_compute_pass() as pass_encoder:
            pass_encoder.push_debug_group("FillInstanceOffsets", (0, 1, 0))
            cursor = self.bind(pass_encoder, self.FillInstanceOffsets.pipeline, self.uniforms)
            pass_encoder.dispatch((num_tiles, 1, 1))
            pass_encoder.pop_debug_group()
        with command_encoder.begin_compute_pass() as pass_encoder:
            pass_encoder.push_debug_group("SpawnInstances", (0, 0, 1))
            cursor = self.bind(pass_encoder, self.SpawnInstances.pipeline, self.uniforms)
            pass_encoder.dispatch((self.strokes.num_strokes, 1, 1))
            pass_encoder.pop_debug_group()
        subres = spy.SubresourceRange()
        subres.layer, subres.layer_count, subres.mip, subres.mip_count = 0, 1, 0, 1
        command_encoder.clear_texture_float(
            texture=self.output_texture,
            range=subres,
            clear_value=spy.math.float4(0, 0, 0, 0)
        )
        with command_encoder.begin_compute_pass() as pass_encoder:
            pass_encoder.push_debug_group("RenderAndBackward", (1, 0, 1))
            cursor = self.bind(pass_encoder, self.RenderAndBackward.pipeline, self.uniforms)
            pass_encoder.dispatch((self.image_width, self.image_height, 1))
            pass_encoder.pop_debug_group()
        with command_encoder.begin_compute_pass() as pass_encoder:
            pass_encoder.push_debug_group("RenderPerStrokeLossAndBackWard", (1, 0, 0))
            cursor = self.bind(pass_encoder, self.RenderPerStrokeLossAndBackWard.pipeline, self.uniforms)
            pass_encoder.dispatch((self.strokes.num_strokes, 1, 1))
            pass_encoder.pop_debug_group()
        with command_encoder.begin_compute_pass() as pass_encoder:
            pass_encoder.push_debug_group("UpdateParameters", (1, 1, 1))
            cursor = self.bind(pass_encoder, self.UpdateParameters.pipeline, self.uniforms)
            pass_encoder.dispatch((self.strokes.num_strokes * 6, 1, 1))
            pass_encoder.pop_debug_group()
        with command_encoder.begin_compute_pass() as pass_encoder:
            pass_encoder.push_debug_group("RegularizeParameters", (1, 1, 0))
            cursor = self.bind(pass_encoder, self.RegularizeParameters.pipeline, self.uniforms)
            pass_encoder.dispatch((self.strokes.num_strokes, 1, 1))
            pass_encoder.pop_debug_group()
    def setup_ui(self):
        screen = self.ui.screen
        window = spy.ui.Window(screen, "Settings", size=spy.float2(500, 300))

        self.fps_text = spy.ui.Text(window, "FPS: 0")
        # self.loss_text = spy.ui.Text(window, "Loss: 0.0")

        def reinit():
            print('reinit')
            device.wait_for_idle()
            self.strokes.random_initialization(self.image_width, self.image_height)

        def capture():
            print('capture')
            device.wait_for_idle()
            output_path = "output.png"
            np_output_texture = self.output_texture.to_numpy()
            np_output_texture = np_output_texture.reshape((self.image_height, self.image_width, 4))
            np_output_texture = pow(np_output_texture, 1.0 / 2.2)
            np_output_texture = (np_output_texture * 255).astype(np.uint8)
            Image.fromarray(np_output_texture).save(output_path)
            print(f"Captured image to {output_path}")

        spy.ui.Button(window, "Reinit", callback=reinit)
        spy.ui.Button(window, "Capture", callback=capture)

        self.adam_eta = spy.ui.SliderFloat(window, "Adam Eta", value=0.005, min=0.001, max=0.01)
        self.stroke_size_var = spy.ui.SliderFloat(window, "Stroke Size Variance", value=20.0, min=1, max=50.0)
        self.standard_stroke_size = spy.ui.SliderFloat(window, "Standard Stroke Size", value=5.0, min=1.0, max=10.0)
        self.lambda_overlap = spy.ui.SliderFloat(window, "Lambda Overlap", value=0.05, min=0.0, max=1.0)

        self.last_mouse_pos = (-1, -1)
    
    def on_keyboard_event(self, event):
        if self.ui.handle_keyboard_event(event):
            return
    
    def on_mouse_event(self, event : spy.MouseEvent):
        if self.ui.handle_mouse_event(event):
            return
        if event.is_move():
            if self.last_mouse_pos != (-1, -1):
                dx_pixels = event.pos.x - self.last_mouse_pos[0]
                dy_pixels = event.pos.y - self.last_mouse_pos[1]

            self.last_mouse_pos = (event.pos.x, event.pos.y)
                

if __name__ == "__main__":
    painter = Painter()
    painter.main()
